'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 1_study_base_graph.py
* @Time: 2025/11/26
* @All Rights Reserve By Brtc
'''
from typing import TypedDict, Annotated

import dotenv
from langchain_openai import ChatOpenAI
from langgraph.constants import START, END
from langgraph.graph import add_messages, StateGraph

dotenv.load_dotenv()

class ChatBotState(TypedDict):
    """聊天消息就是状态， 消息是一个字典， add_messages 就是一个归纳函数"""
    messages:Annotated[list, add_messages]
    node_name:str

llm = ChatOpenAI(model="gpt-4o-mini")

def chat_bot(state: ChatBotState)->ChatBotState:
    """聊天机器人节点函数"""
    print(state["messages"])
    ai_message = llm.invoke(state["messages"])
    return {"messages": [ai_message], "node_name": "chat_bot"}

#1、构建图架构
graph_builder = StateGraph(ChatBotState)
#2、添加聊天节点
graph_builder.add_node("llm", chat_bot)
#3、通过边连接节点  star ----> llm ------> end
graph_builder.add_edge(START, "llm")
graph_builder.add_edge("llm", END)
#4、编译图架构
graph_bin = graph_builder.compile()
#5、运行图架构
"""  
{
"messages":[("human","你好你是？")],
"node_name":"star_bot"
}
----->start --------> llm -------->end
{
'messages': [HumanMessage(content='你好你是？', additional_kwargs={}, response_metadata={}, id='848c9fbd-42ee-4703-88eb-c52c1a9922d3'), AIMessage(content='你好！我是一个人工智能助手，旨在回答你的问题和提供帮助。请问有什么我可以为你做的？', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 28, 'prompt_tokens': 11, 'total_tokens': 39, 'completion_tokens_details': {'accepted_prediction_tokens': 0, 'audio_tokens': 0, 'reasoning_tokens': 0, 'rejected_prediction_tokens': 0}, 'prompt_tokens_details': {'audio_tokens': 0, 'cached_tokens': 0}}, 'model_name': 'gpt-4o-mini-2024-07-18', 'system_fingerprint': 'fp_efad92c60b', 'id': 'chatcmpl-Cg8cn2NHXBIbVVIcdKEbI3CMPMSrC', 'service_tier': None, 'finish_reason': 'stop', 'logprobs': None}, id='run--7752e717-d2ca-4b90-bf43-6c75860904c4-0', usage_metadata={'input_tokens': 11, 'output_tokens': 28, 'total_tokens': 39, 'input_token_details': {'audio': 0, 'cache_read': 0}, 'output_token_details': {'audio': 0, 'reasoning': 0}})], 
'node_name': 'chat_bot'
}
"""
print(graph_bin.invoke({"messages":[("human","你好你是？")], "node_name":"star_bot"}))
